package com.pj.controller;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.pj.conf.RagflowConfig;
import com.pj.service.RagFlowChatService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import reactor.core.Disposable;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author zzy
 * @date 2025/6/21 16:48
 * @description : 会话管理
 */
@Controller
@RequestMapping("/ragFlow/chat")
public class RagFlowChatController {

    @Autowired
    private RagFlowChatService ragFlowChatService;

    private final WebClient webClient;
    private final RagflowConfig ragflowConfig;

    public RagFlowChatController(WebClient.Builder webClient, RagflowConfig ragflowConfig) {
        this.webClient = webClient.baseUrl(ragflowConfig.getBaseUrl()).build();
        this.ragflowConfig = ragflowConfig;
    }

    /**
     * 创建会话
     * chatId 聊天助手Id
     * 会话的名称
     */
    @ResponseBody
    @RequestMapping("/add")
    public List<Map<String , Object>> addChat(String chatId,String name){
        if(chatId == null || chatId.isEmpty()){
           throw new  RuntimeException("请输出chatId");
        }
        if(name == null || name.isEmpty()){
            throw new  RuntimeException("请输出chatId");
        }
        String ragFlowApi = "/api/v1/chats/"+chatId+"/sessions";
        // 构造请求体
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("name", name);  // 会话名称
//        requestBody.put("user_id", "custom_user_id"); // 可选，用户自定义 ID（如果不需要可以不加）

        return ragFlowChatService.addChat(ragFlowApi,requestBody);

    }

    /**
     * 与聊天会话交谈
     */
    @ResponseBody
    @RequestMapping("/chat")
    public Map<String, Object> chat(String question, boolean stream,String sessionId,String chatId) {
        if (chatId == null || chatId.isEmpty()){
            throw new RuntimeException("chatId为空");
        }
        if (question == null || question.isEmpty()){
            throw new RuntimeException("question");
        }
        String ragFlowApi = "/api/v1/chats/" + chatId + "/completions";

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("question", question); // 必填
        requestBody.put("stream", stream);
        requestBody.put("session_id", sessionId);
//        requestBody.put("user_id", "custom_user_id");

        return ragFlowChatService.chat(ragFlowApi, requestBody);
    }

    /**
     * 与聊天会话交谈
     * 使用流式输出
     */
    @RequestMapping(value = "/chatStream", produces = "text/event-stream")
    public ResponseBodyEmitter chatStream(String question, String sessionId, String chatId, HttpServletResponse response) {

        response.setContentType("text/event-stream;charset=UTF-8");
        response.setHeader("Cache-Control", "no-cache");
        response.setHeader("Connection", "keep-alive");

        if (chatId == null || chatId.isEmpty()) {
            throw new RuntimeException("chatId 为空");
        }
        if (sessionId == null || sessionId.isEmpty()) {
            throw new RuntimeException("sessionId 为空");
        }
        if (question == null || question.isEmpty()) {
            throw new RuntimeException("question 为空");
        }

        System.out.println("chatId: " + chatId);
        System.out.println("sessionId: " + sessionId);

        ResponseBodyEmitter emitter = new ResponseBodyEmitter();

        try {
            // 立即发送心跳，防止前端报错
            emitter.send("data: [连接已建立]\n\n");
        } catch (IOException e) {
            e.printStackTrace();
        }

        String ragFlowApi = ragflowConfig.getBaseUrl() + "/api/v1/chats/" + chatId + "/completions";

        Map<String, Object> queryMap = new HashMap<>();
        queryMap.put("question", question);
        queryMap.put("stream", true); // 强制写死流式，防止前端误传
        queryMap.put("session_id", sessionId);

        AtomicBoolean completed = new AtomicBoolean(false);

        WebClient.create()
                .post()
                .uri(ragFlowApi)
                .header("Authorization", "Bearer " + ragflowConfig.getApiKey())
                .bodyValue(queryMap)
                .retrieve()
                .bodyToFlux(String.class)
                .doOnNext(data -> {
                    if (completed.get()) return;

                    try {
                        System.out.println("流式数据：" + data);

                        JSONObject jsonObject = JSON.parseObject(data);
                        Integer code = jsonObject.getInteger("code");
                        Object dataObj = jsonObject.get("data");
                        String message = jsonObject.getString("message");

                        // 失败逻辑，code不为0时，认为请求失败，提示前端并关闭流
                        if (code == null || code != 0) {
                            if (completed.compareAndSet(false, true)) {
                                emitter.send("data: 错误: " + (message != null ? message : "未知错误") + "\n\n");
                                emitter.completeWithError(new RuntimeException("接口返回错误，code=" + code));
                            }
                            return;
                        }

                        // 正常回答，dataObj是Map且包含answer字段
                        if (dataObj instanceof Map) {
                            Map<String, Object> dataMap = (Map<String, Object>) dataObj;
                            Object answerObj = dataMap.get("answer");
                            if (answerObj != null) {
                                String answer = answerObj.toString().replace("\n", "\\n");
                                emitter.send("data: " + answer + "\n\n");
                            }
                        }
                        // 结束标志，dataObj为true，表示流结束
                        else if (dataObj instanceof Boolean && (Boolean) dataObj) {
                            if (completed.compareAndSet(false, true)) {
                                emitter.complete();
                            }
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                        if (completed.compareAndSet(false, true)) {
                            try {
                                emitter.send("data: 流式解析异常: " + e.getMessage() + "\n\n");
                            } catch (IOException ioException) {
                                ioException.printStackTrace();
                            }
                            emitter.completeWithError(e);
                        }
                    }
                })
                .doOnError(error -> {
                    if (completed.compareAndSet(false, true)) { // ✅ 添加原子检查
                        try {
                            emitter.send("data: 发生错误: " + error.getMessage() + "\n\n");
                            emitter.completeWithError(error);
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    }
                })

                .doOnComplete(() -> {
                    if (completed.compareAndSet(false, true)) { // ✅ 添加原子检查，防止重复关闭
                        try {
                            emitter.complete();
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    }
                })

                .subscribe();

        return emitter;
    }



    /**
     * 跳转到聊天页面
     */
    @GetMapping("/chatPage")
    public String chatPage() {
        return "chat";
    }

    /**
     * 列出聊天助手的会话列表
     */
    @ResponseBody
    @RequestMapping("/list")
    public Map<String, Object> listChat(String chatId,
                                        @RequestParam(required = false) String page,
                                        @RequestParam(required = false) String pageSize,
                                        @RequestParam(required = false) String orderBy,
                                        @RequestParam(required = false) String desc,
                                        @RequestParam(required = false) String sessionName,
                                        @RequestParam(required = false) String sessionId,
                                        @RequestParam(required = false) String userId) {

        if (chatId == null || chatId.trim().isEmpty()) {
            throw new RuntimeException("chatId 不能为空");
        }
        // 基础 URL
        StringBuilder ragFlowAPI = new StringBuilder(String.format("/api/v1/chats/%s/sessions", chatId));

        // 动态参数拼接
        List<String> params = new ArrayList<>();
        if (page != null && !page.trim().isEmpty()) {
            params.add("page=" + page);
        }
        if (pageSize != null && !pageSize.trim().isEmpty()) {
            params.add("page_size=" + pageSize);
        }
        if (orderBy != null && !orderBy.trim().isEmpty()) {
            params.add("orderby=" + orderBy);
        }
        if (desc != null && !desc.trim().isEmpty()) {
            params.add("desc=" + desc);
        }
        if (sessionName != null && !sessionName.trim().isEmpty()) {
            params.add("name=" + sessionName);
        }
        if (sessionId != null && !sessionId.trim().isEmpty()) {
            params.add("id=" + sessionId);
        }
        if (userId != null && !userId.trim().isEmpty()) {
            params.add("user_id=" + userId);
        }

        // 拼接查询参数
        if (!params.isEmpty()) {
            ragFlowAPI.append("?").append(String.join("&", params));
        }

        // 调用实现层
        List<Map<String, Object>> sessions = ragFlowChatService.listChat(ragFlowAPI.toString());

        // 封装返回结果
        Map<String, Object> result = new HashMap<>();
        result.put("code", 0);
        result.put("data", sessions);

        return result;
    }




}
